'''
Created on Oct 20, 2009

@author: mkiyer
'''

from optparse import OptionParser
import sys
import os
import glob
import logging

from veggie.app.rnaseq.cufflinks import CuffCompare
import veggie.db.sample as sdb
from veggie.db.sample.samplegroup import parse_samplegroups_xml
import veggie.genome.genome as genome

def fetch_libraries_from_sample(sample_name, sdbi, qc_filter=True):
    logger = logging.getLogger(__name__)    
    # get the libraries associated with this sample
    libraries = sdbi.get_libraries_by_sample_name(sample_name, best=True)
    for library in libraries:
        if library.qc_status == False:
            logger.warning('library %s: marked QC-FAIL' % library.id)
            if qc_filter:
                logger.warning('library %s: excluded from query' % library.id)
                continue
        yield library.id

def extract_chroms_from_gtf(gtf_file, suffix, genome_version):
    chroms = genome.get_genome(genome_version).get_chrom_names()  
    chrom_files = {}
    for chrom in chroms:
        out_file = '_'.join([chrom, suffix]) + '.gtf'
        chrom_files[chrom] = open(out_file, 'w')
    for line in open(gtf_file):
        chrom = line.strip().split('\t')[0]
        chrom_files[chrom].write(line)
    for chrom in chroms:
        chrom_files[chrom].close()

if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    optionparser = OptionParser("usage: %prog [options] <rnaseq_results_path> <samples.xml>")
    optionparser.add_option("-o", "--output-file",
                            dest="output_file",
                            default=None)
    optionparser.add_option("-r", "--refgtf",
                            dest="ref_gtf",
                            default=None)
    optionparser.add_option("--chrom", dest="chrom", action="store_true", default=False)
    (options, args) = optionparser.parse_args()
    
    rnaseq_results_path = args[0]
    sample_xml_file = args[1]
    sdbi_path = '/lab/mkiyer/refdb/sample/sampledb_2010_Feb_04.pickle'

    # load sampledb snapshot
    sdbi = sdb.get_sampledb_instance(sdbi_path)
    sample_groups = list(parse_samplegroups_xml(sample_xml_file))
    library_ids = set()
    for sgroup in sample_groups:
        for sample_name in sgroup.samples:
            library_ids.update(list(fetch_libraries_from_sample(sample_name, sdbi, qc_filter=True)))
            
    input_files = []
    for filename in glob.iglob(os.path.join(rnaseq_results_path, '*')):
        if os.path.isdir(filename) and (os.path.basename(filename) in library_ids):
            gtf_file = os.path.join(rnaseq_results_path, filename, 'transcripts.gtf')
            if options.chrom:
                name = os.path.basename(os.path.dirname(gtf_file))
                extract_chroms_from_gtf(gtf_file, name, 'hg18')

    sys.exit(0)
    input_files.append(gtf_file)
    
    executable = '/lab/sw/bioinfo/rnaseq/cufflinks/cufflinks-0.8.2.0.Linux_x86_64/cuffcompare'
    opts = '-o results -r /lab/mkiyer/refdb/annotations/custom/annotated_genes_2010_Jan_12.gtf -R'
    
    command = ' '.join([executable, opts, ' '.join(input_files)])
    print command
    os.system(command)
